#include<linux/module.h>
#include<linux/printk.h>
#include<linux/kobject.h>
#include<linux/kernel.h>
#include<linux/string.h>
#include "memcg.h"

static struct mem_cgroup *(* _mem_cgroup_iter)(struct mem_cgroup *,struct mem_cgroup *,struct mem_cgroup_reclaim_cookie *);
static void (*_mem_cgroup_iter_break)(struct mem_cgroup*, struct mem_cgroup*);
static struct zone *(*_next_zone)(struct zone *zone);
static struct address_space *(*_page_mapping)(struct page *page);
struct pglist_data *(*_first_online_pgdat)(void);

#define for_each_mem_cgroup(iter, start)           \
    for (iter = _mem_cgroup_iter(start, NULL, NULL);  \
         iter != NULL;              \
         iter = _mem_cgroup_iter(start, iter, NULL))
#define _for_each_zone(zone)                 \
    for (zone = (_first_online_pgdat())->node_zones; \
         zone;                  \
         zone = _next_zone(zone))

static void free_item(void *args)
{
    struct inode_item *item = (struct inode_item *)args;
    kfree(item->filename);
    kfree(item);
}

static void memcg_free_item(struct memcg_item *cgitem)
{
   radix_delete_all(&cgitem->inode_root, free_item);
}


static int get_page_inode(struct memcg_item *cgitem, struct page *page)
{
    struct address_space *mapping;
    struct inode *inode;
    struct inode_item *item;
    unsigned long cached;
    int deleted;
    int ret = 0;
    char *fileName;

    fileName = kzalloc(PATH_MAX, GFP_ATOMIC);
    if (!fileName)
        return 0;
    get_page(page);
    if (PageAnon(page)) {
        cgitem->anon++;
        goto _skip;
    }
    if (!_page_is_file_cache(page))
       cgitem->shmem++;
    else
        cgitem->file++;
#ifdef LINUX_310
    mapping = _page_mapping(page);
#else
    mapping = page_mapping(page);
#endif
    if (!mapping)
        goto _skip;
    inode = mapping->host;
    if (!inode)
        goto _skip;

    item = radix_lookup(&cgitem->inode_root, (unsigned long)inode);
    if (item) {
        item->nr_pages++;
        goto _skip;
    }
    scan_inode_name(inode, fileName, PATH_MAX, &cached, &deleted);
    if (!cached)
        goto _skip;

    item = kzalloc(sizeof(*item), GFP_ATOMIC);
    if (!item)
        goto _skip;
    item->i_ino = inode->i_ino;
    item->nr_pages = 1;
    item->filename = kmemdup(fileName, strlen(fileName) + 1, GFP_ATOMIC);
    item->cached = cached;
    item->deleted = deleted;
    item->shmem = page_is_shmem(page);
    item->node.key = (unsigned long)inode;
    list_add_tail(&item->inode, &cgitem->head);
    ret = radix_insert(&cgitem->inode_root, (unsigned long)inode, (void*)item);
    if(ret)
        pr_info("insert file:%s error\n", item->filename);
    cgitem->num_file++;
_skip:
    kfree(fileName);
    put_page(page);
    return 0;
}


static void get_lru_page(struct memcg_item *item, struct lruvec *vec)
{
    struct page *page, *tmp;
    enum lru_list lru;
    unsigned long flags;

#ifdef LINUX_310
    struct zone *lruzone;
    lruzone = vec->zone;
    if (!lruzone) {
        pr_err("lru zone error for memcg:%px cg:%s\n", item->memcg, item->cgname);
        return;
    }
    if(lruzone)
        spin_lock_irqsave(&lruzone->lru_lock, flags);
#else
    struct pglist_data *pgdat;
    pgdat = vec->pgdat;
    if (!pgdat) {
        pr_err("lru pgdata error for memcg:%px cg:%s\n", item->memcg, item->cgname);
        return;
    }
    if(pgdat)
        spin_lock_irqsave(&pgdat->lru_lock, flags);
#endif
    for_each_lru(lru) {
        struct list_head *list = &vec->lists[lru];
        list_for_each_entry_safe(page, tmp, list, lru) {
            get_page_inode(item, page);
        }
    }
#ifdef LINUX_310
   if (lruzone)
       spin_unlock_irqrestore(&lruzone->lru_lock, flags);
#else
    if (pgdat)
        spin_unlock_irqrestore(&pgdat->lru_lock, flags);
#endif
}

static void get_memcg_page(struct memcg_item *item)
{
#ifdef LINUX_310
    struct mem_cgroup_per_zone *mz;
    struct zone *z;
    _for_each_zone(z) {
	    if((unsigned)zone_to_nid(z) >= nr_node_ids)
		    continue;
	    mz = &item->memcg->info.nodeinfo[zone_to_nid(z)]->zoneinfo[zone_idx(z)];
        get_lru_page(item, &mz->lruvec);
    }
#else
    struct mem_cgroup_per_node *mz;
    int nid;
    for_each_node(nid) {
        mz = mem_cgroup_nodeinfo(item->memcg, nid);
        get_lru_page(item, &mz->lruvec);
    }
#endif
}

static void memcg_get_name(struct mem_cgroup *memcg, char *name, unsigned int len)
{
    char *end;
    int pos;
    struct cgroup *cg = memcg->css.cgroup;
#ifdef LINUX_310
    if (!cg)
        return;
    rcu_read_lock();
    cgroup_path(cg, name, PATH_MAX);
    rcu_read_unlock();
    end = name+strlen("/sys/fs/cgroup/memory");
    memmove(end, name, strlen(name)+1);
    prepend(&end, &len, "/sys/fs/cgroup/memory", strlen("/sys/fs/cgroup/memory"));
    pr_err("cg:name: %s, len:%d\n",name,strlen(name));
#else
    struct kernfs_node *kn;
    struct kernfs_node *pkn;
    if (!cg|| !cg->kn)
        return;
    kn = cg->kn;

    kernfs_get(kn);
    end = name + len - 1;
    prepend(&end, &len, "\0", 1);
    pkn = kn;
    while (pkn) {
        pos = prepend(&end, &len, pkn->name, strlen(pkn->name));
        if (pos)
            break;
        if ((pkn == pkn->parent) || !pkn->parent)
            break;
        pos = prepend(&end, &len, "/", 1);
        if (pos)
            break;
        pkn = pkn->parent;
    }

    prepend(&end, &len, "/sys/fs/cgroup/memory", strlen("/sys/fs/cgroup/memory"));

    kernfs_put(kn);
    memmove(name, end, strlen(end) + 1);
#endif
}

static struct memcg_item *memcg_init_item(struct mem_cgroup *cg)
{
    struct memcg_item *item = NULL;
    char *fileName;

    fileName = kzalloc(PATH_MAX, GFP_ATOMIC);
    if (!fileName)
        return NULL;
    item = kzalloc(sizeof(*item), GFP_ATOMIC);
    if (!item) {
        goto _out;
    }
    memcg_get_name(cg, fileName, PATH_MAX);
    item->memcg = cg;
#ifdef LINUX_310
    item->size = cg->res.usage;
    item->cgname = kmemdup(fileName, strlen(fileName) + 2, GFP_ATOMIC);
    if(item->cgname)
        item->cgname[strlen(fileName) + 1] = '\0';
#else
    item->size = page_counter_read(&cg->memory);
    item->cgname = kmemdup_nul(fileName, strlen(fileName) + 1, GFP_ATOMIC);
#endif
    INIT_LIST_HEAD(&item->head);
    INIT_LIST_HEAD(&item->offline);
    radix_init(&item->inode_root);
_out:
    kfree(fileName);
    return item;
}

int memcg_dump_to_user(struct memcg_info *info, struct memcg_info_user __user *result)
{
    struct memcg_item *cgitem;
    struct inode_item *item;
    struct memcg_info_user res;
    struct memcg_item_user *cgitem_u, *tmp, *user_head;
    struct memcg_item_user tmp_memcg;
    struct inode_item_user *tmp2, *free_tmp;
    int count_m = 0, i = 0, j = 0, count_i;
    int ret = 0;

    if(copy_from_user(&res, result, sizeof(res)))
    {
        pr_err("can not copy from user %d:%d\n",count_m,__LINE__);
        ret = copy_to_user(result, &count_m, sizeof(count_m));
        return 0;
    }

    if (!res.nr || !res.items)
    {
        pr_err("num %d ,items %p \n", res.nr, res.items);
        ret = copy_to_user(result, &count_m, sizeof(count_m));
        return 0;
    }

    i = res.nr > info->nr ? info->nr : res.nr;
    if(i == 0)
    {
        ret = copy_to_user(result, &count_m, sizeof(count_m));
        return 0;
    }
    if((tmp = vzalloc(sizeof(struct memcg_item_user) * i)) == NULL)
    {
        pr_err("vmalloc error %d:%d\n",count_m,__LINE__);
        ret = copy_to_user(result, &count_m, sizeof(count_m));
        return 0;
    }

    cgitem_u = res.items;
    user_head = tmp;

    list_for_each_entry(cgitem, &info->head, offline) {
        if(count_m >= i)
            break;
        pr_err("cg:%s memory:%lu file:%lu anon:%lu shmem:%lu num_file:%d\n", cgitem->cgname, cgitem->size, cgitem->file, cgitem->anon, cgitem->shmem, cgitem->num_file);
        strncpy(tmp->cgname, cgitem->cgname,strlen(cgitem->cgname));
        tmp->size = cgitem->size;
        tmp->file = cgitem->file;
        tmp->anon = cgitem->anon;
        tmp->shmem = cgitem->shmem;
        tmp->num_file = cgitem->num_file;

        j = INODE_LIMIT > tmp->num_file ? tmp->num_file : INODE_LIMIT;
        if(j <= 0)
        {
            if(copy_from_user(&tmp_memcg, cgitem_u, sizeof(tmp_memcg)))
            {
                pr_err("can not copy from user %d:%d\n",count_m,__LINE__);
                ret = copy_to_user(result, &count_m, sizeof(count_m));
                vfree(tmp->inode_items);
                vfree(user_head);
                return 0;
            }
            tmp->inode_items = tmp_memcg.inode_items;
            copy_to_user(cgitem_u, tmp, sizeof(struct memcg_item_user));
            tmp++;
            cgitem_u++;
            count_m++;
            continue;
        }
        if((tmp2 = vzalloc(sizeof(struct inode_item_user) * j)) == NULL)
        {
            pr_err("vmalloc error %d:%d\n",count_m,__LINE__);
            ret = copy_to_user(result, &count_m, sizeof(count_m));
            vfree(user_head);
            return 0;
        }
        free_tmp = tmp2;
        count_i = 0;

        list_for_each_entry(item, &cgitem->head,inode) {
            if(count_i >= j)
               break;
            pr_err("ino:%lu, filename:%s cached:%lu nr_pages:%d deleted:%d shmem:%d\n", item->i_ino, item->filename, item->cached, item->nr_pages, item->deleted, item->shmem);
            tmp2->i_ino = item->i_ino;
            strcpy(tmp2->filename, item->filename);
            tmp2->cached = item->cached;
            tmp2->nr_pages = item->nr_pages;
            tmp2->deleted = item->deleted;
            tmp2->shmem = item->shmem;
            tmp2++;
            count_i++;
        }
        if(copy_from_user(&tmp_memcg, cgitem_u, sizeof(tmp_memcg)))
        {
            pr_err("can not copy from user %d:%d\n",count_m,__LINE__);
            ret = copy_to_user(result, &count_m, sizeof(count_m));
            vfree(tmp->inode_items);
            vfree(user_head);
            return 0;
        }
        copy_to_user(tmp_memcg.inode_items, free_tmp, sizeof(struct inode_item_user) * j);
        tmp->inode_items = tmp_memcg.inode_items;
        copy_to_user(cgitem_u, tmp, sizeof(struct memcg_item_user));
        vfree(free_tmp);
        tmp++;
        cgitem_u++;
        count_m++;
    }
    i = copy_to_user(result, &count_m, sizeof(count_m));
    vfree(user_head);
    return i;

}

void memcg_dump(struct memcg_info *info)
{
    struct memcg_item *cgitem;
    struct inode_item *item;
    int a = 0;
    list_for_each_entry(cgitem, &info->head, offline) {
        pr_err("cg:%s memory:%lu file:%lu anon:%lu shmem:%lu num_file:%d\n", cgitem->cgname, cgitem->size, cgitem->file, cgitem->anon, cgitem->shmem, cgitem->num_file);
        list_for_each_entry(item, &cgitem->head,inode) {
            pr_err("ino:%lu, filename:%s cached:%lu nr_pages:%d deleted:%d shmem:%d\n", item->i_ino, item->filename, item->cached, item->nr_pages, item->deleted, item->shmem);
            a = 1;
        }
    }
}

void memcg_free_all(struct memcg_info *info)
{
    struct memcg_item *cgitem;
    struct memcg_item *tmp;

    list_for_each_entry_safe(cgitem, tmp, &info->head, offline) {
        memcg_free_item(cgitem);
        kfree(cgitem->cgname);
        kfree(cgitem);
    }
}

static int memcg_init(void)
{

    if (_mem_cgroup_iter && _mem_cgroup_iter_break)
        return 0;

    _mem_cgroup_iter = kallsyms_lookup_name("mem_cgroup_iter");
    if (!_mem_cgroup_iter) {
        pr_err("lookup mem cgroup iter error\n");
        return -1;
    }

    _mem_cgroup_iter_break = kallsyms_lookup_name("mem_cgroup_iter_break");
    if (!_mem_cgroup_iter_break) {
        pr_err("lookup iter break error\n");
        return -1;
    }
    _next_zone = kallsyms_lookup_name("next_zone");
    if (!_next_zone) {
        pr_err("next_zone error\n");
        return -1;
    }

    _first_online_pgdat = kallsyms_lookup_name("first_online_pgdat");
    if (!_first_online_pgdat) {
        pr_err("first_online_pgdat error\n");
        return -1;
    }
    _page_mapping = kallsyms_lookup_name("page_mapping");
    if (!_page_mapping) {
        pr_err("page_mapping error\n");
        return -1;
    }

    return 0;
}

int memcg_scan(struct memcg_info *info, struct mem_cgroup *start, int offline)
{
    struct mem_cgroup *iter = NULL;
    struct memcg_item *item;

    if (memcg_init())
        return -1;

    for_each_mem_cgroup(iter, start) {
#ifdef LINUX_310
        if (offline && !(iter->css.flags&CSS_DYING))
#else
        if (offline && (iter->css.flags&CSS_ONLINE))
#endif
            continue;
        item =  memcg_init_item(iter);
        if (!item)
            continue;
        get_memcg_page(item);
        info->nr++;
        list_add_tail(&item->offline, &info->head);
    }
	return 0;
}

struct mem_cgroup *memcg_get_by_name(char *cgname)
{
    struct mem_cgroup *iter = NULL;
    char *fileName;

    if (memcg_init())
        return 0;
    fileName = kzalloc(PATH_MAX, GFP_ATOMIC);
    if (!fileName) {
        pr_err("alloc memory failed:%pF\n", __FUNCTION__);
        return 0;
    }

    for_each_mem_cgroup(iter, NULL) {
         memcg_get_name(iter, fileName, PATH_MAX);
        if (!strncmp(fileName, cgname, strlen(cgname))) {
            pr_err("filename:%s, cgname:%s\n",fileName, cgname);
            _mem_cgroup_iter_break(NULL, iter);
            break;
        }
    }

    kfree(fileName);
    return iter;
}
int memcg_dying_main(unsigned long arg)
{
    pr_err("try_to dump to user\n");
    struct memcg_info info;
    struct mem_cgroup *memcg = NULL;
    info.nr = 0;
    char cgname_u[NAME_LEN];
    struct memcg_info_user *res = (struct memcg_info_user*)arg;

    memset(cgname_u, 0, sizeof(cgname_u));
    copy_from_user(cgname_u, res->cgname, NAME_LEN);
    if(strlen(cgname_u) > 0)
    {
        if(cgname_u[strlen(cgname_u)-1] == '/')
            cgname_u[strlen(cgname_u)-1] = '\0';
        pr_err("cgname:%s\n",cgname_u);
        memcg = memcg_get_by_name(cgname_u);
    }

    INIT_LIST_HEAD(&info.head);
    memcg_scan(&info, memcg, 1);
    memcg_dump_to_user(&info, (struct memcg_info_user*)arg);
    memcg_free_all(&info);

    return 0;
}
int memcg_one_main(unsigned long arg)
{
    char *cgname = "/sys/fs/cgroup/memory/agent";
    struct mem_cgroup *memcg = NULL;
    struct memcg_info info;
    char cgname_u[NAME_LEN];
    struct memcg_info_user *res = (struct memcg_info_user*)arg;
    info.nr = 0;

    memset(cgname_u, 0, sizeof(cgname_u));
    copy_from_user(cgname_u, res->cgname, NAME_LEN);
    pr_err("cgname:%s\n",cgname_u);
    if(strlen(cgname_u) > 0)
        if(cgname_u[strlen(cgname_u)-1] == '/')
            cgname_u[strlen(cgname_u)-1] = '\0';
    INIT_LIST_HEAD(&info.head);
    memcg = memcg_get_by_name(cgname_u);
    if(memcg != NULL)
        memcg_scan(&info, memcg, 0);
    memcg_dump_to_user(&info, (struct memcg_info_user*)arg);
    memcg_free_all(&info);
    return 0;
}
